package com.open.web.filter;

import com.open.basic.constant.TraceId;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.UUID;


/**
 * Trace跟踪处理器
 */
@Slf4j
@Component
public class TraceParentHandler implements HandlerInterceptor {

    @Value("${logging.trace.enabled:true}")
    private boolean enabled = true;
    private static final String HEADER_TRACE_PARENT = "traceparent";
    public static final String TRACE_ID = TraceId.TRACE_ID;

    /**
     * 清除MDC中的traceId
     * @param request
     * @param response
     * @param handler
     * @param ex
     * @throws Exception
     */
    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response,
                                Object handler, Exception ex) throws Exception {
        MDC.remove(TRACE_ID);
        HandlerInterceptor.super.afterCompletion(request, response, handler, ex);
    }

    /**
     * MDC增加traceId
     * @param request
     * @param response
     * @param handler
     * @return
     * @throws Exception
     */
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response,
                             Object handler) throws Exception {
        if (enabled) {
            //开始获取header中的字段，如果没有生成一个
            String traceId = getTraceId(request);
            MDC.put(TRACE_ID, traceId);
            return true;
        } else {
            return HandlerInterceptor.super.preHandle(request, response, handler);
        }
    }

    /**
     * traceId生成
     * @param request
     * @return
     */
    private String getTraceId(HttpServletRequest request) {
        String traceParent = request.getHeader(HEADER_TRACE_PARENT);
        if (StringUtils.hasLength(traceParent)) {
            String[] tracesInfos = traceParent.split("-");
            return tracesInfos[1];
        } else {
            String traceId = request.getHeader(TRACE_ID);
            if (StringUtils.hasLength(traceId)) {
                return traceId;
            }
            //生成一个bing设置
            return UUID.randomUUID().toString().replace("-", "");
        }
    }
}
